import os
from pascal_voc_io import PascalVocReader
import numpy as np

root_path = "/disk1/DataSet/widerface/"
anno_path = root_path + "wider_face_split/wider_face_train_bbx_kp_gt.txt"
with open(anno_path, 'r') as file:
    annos = file.readlines()


image_names = []
for anno in annos:
    anno = str(anno).rstrip()
    if anno.endswith(".jpg"):
        image_names.append(anno)


label_save = open(root_path + "wider_face_split/wider_face_train_bbx_kp_gt_xml.txt", 'w')
for image_name in image_names:
    label_name = image_name[:-4] + '.xml'
    label_path = os.path.join(root_path, 'WIDER_train', 'annotations', label_name)
    if not os.path.isfile(label_path):
        print(label_path, 'is not valid xml')
        continue
    tVocParseReader = PascalVocReader(label_path)
    shapes = tVocParseReader.getShapes()

    if (len(shapes) > 0):
        label_save.write("{}\n".format(image_name))
        label_save.write("{}\n".format(len(shapes)))

    for idx, shape in enumerate(shapes):
        label, bndbox, landmarks, gender, mask, glass, hat, cellphone, smoking, yawpose, blur, expression, illumination, occlusion, pose, invalid = shape
        x1 = bndbox[0]
        y1 = bndbox[1]
        w  = bndbox[2] - bndbox[0]
        h  = bndbox[3] - bndbox[1]
        lm_x1 = 0
        lm_y1 = 0
        lm_x2 = 0
        lm_y2 = 0
        lm_x3 = 0
        lm_y3 = 0
        lm_x4 = 0
        lm_y4 = 0
        lm_x5 = 0
        lm_y5 = 0
        if landmarks[0][0] > 0:
            lm_x1 = landmarks[0][0]
            lm_y1 = landmarks[0][1]
            lm_x2 = landmarks[1][0]
            lm_y2 = landmarks[1][1]
            lm_x3 = landmarks[2][0]
            lm_y3 = landmarks[2][1]
            lm_x4 = landmarks[3][0]
            lm_y4 = landmarks[3][1]
            lm_x5 = landmarks[4][0]
            lm_y5 = landmarks[4][1]
        label_save.write("%d %d %d %d %d %d %d %d %d %d %d %d %d %d\n"%(x1, y1, w, h, lm_x1, lm_y1, lm_x2, lm_y2, lm_x3, lm_y3, lm_x4, lm_y4, lm_x5, lm_y5))
label_save.close()